/*
 *         (C) COPYRIGHT NBI Limited
 *              ALL RIGHT RESERVED
 *
 * File        : fastflow/graph/node.cc
 * Authors     : dzhang
 * Create Time : 2021/09/04 18:10:42
 * Description :
 *
 */

#include "fastflow/graph/node.h"

#include <iostream>
#include <vector>
#include <numeric>
#include <string>
#include <functional>

#include "fastflow/gsl/gsl"
#include "fastflow/graph/constants.h"
#include "fastflow/graph/graph.h"

using namespace ONNX_NAMESPACE;

namespace fastflow {

Node::EdgeEnd::EdgeEnd(const Node& node, const NodeArg& node_arg) noexcept : node_(&node), node_arg_(&node_arg) {}

Node::EdgeEnd::EdgeEnd(const Node& node) noexcept : node_(&node), node_arg_(nullptr) {}

const Node& Node::EdgeEnd::GetNode() const noexcept {
    return *node_;
}

const NodeArg* Node::EdgeEnd::GetNodeArg() const noexcept {
    return node_arg_;
}

Node::NodeConstIterator::NodeConstIterator(EdgeConstIterator p_iter) {
    m_iter = p_iter;
}

bool Node::NodeConstIterator::operator==(const NodeConstIterator& p_other) const {
    return m_iter == p_other.m_iter;
}

bool Node::NodeConstIterator::operator!=(const NodeConstIterator& p_other) const {
    return m_iter != p_other.m_iter;
}

void Node::NodeConstIterator::operator++() {
    ++m_iter;
}

void Node::NodeConstIterator::operator--() {
    --m_iter;
}

const Node& Node::NodeConstIterator::operator*() {
    return (*m_iter).GetNode();
}

NodeIndex Node::Index() const noexcept {
    return index_;
}

const std::string& Node::Name() const noexcept {
    return name_;
}

const std::string& Node::OpType() const noexcept {
    return op_type_;
}

const std::string& Node::Description() const noexcept {
    return description_;
}

const std::string& Node::Domain() const noexcept {
    return domain_;
}

const OpSchema* Node::Op() const noexcept {
    return op_;
}

Node::Type Node::NodeType() const noexcept {
    return node_type_;
}

void Node::SetNodeType(Node::Type node_type) noexcept {
    node_type_ = node_type;
}

const Function* Node::GetFunctionBody() const noexcept {
    return func_body_;
}

void Node::SetFunctionBody(const Function& func) {
    func_body_ = &func;
    op_ = &func.OpSchema();
}

const std::string& Node::GetExecutionProviderType() const noexcept {
    return execution_provider_type_;
}

void Node::SetExecutionProviderType(ProviderType execution_provider_type) {
    execution_provider_type_ = execution_provider_type;
}

void Node::ToProto(NodeProto& proto) const {
    // Set name.
    proto.set_name(name_);
    // Set op type.
    proto.set_op_type(op_type_);
    // Set op domain;
    proto.set_domain(domain_);
    // Set doc string.
    proto.set_doc_string(description_);

    // Set attributes.
    proto.clear_attribute();
    for (auto attribute : attributes_) {
        const gsl::not_null<AttributeProto*> attr{proto.add_attribute()};
        *attr = attribute.second;
    }

    // Set inputs' definitions.
    proto.clear_input();
    for (auto& input_def : definitions_.input_defs) {
        proto.add_input(input_def->Name());
    }

    // Set outputs' definitions.
    proto.clear_output();
    for (auto& output_def : definitions_.output_defs) {
        proto.add_output(output_def->Name());
    }
}

void Node::Init(const std::string& name, const std::string& op_type, const std::string& description,
                const std::vector<NodeArg*>& input_args, const std::vector<NodeArg*>& output_args,
                const NodeAttributes* attributes, const std::string& domain) {
    name_ = name;
    op_type_ = op_type;
    description_ = description;
    definitions_.input_defs = input_args;
    definitions_.output_defs = output_args;
    domain_ = domain;
    if (kOnnxDomainAlias == domain_) {
        domain_ = kOnnxDomain;
    }

    // Set each arg count as 1 by default.
    // It could be adjusted when resolving the node with its operator
    // information.
    definitions_.input_arg_count.assign(input_args.size(), 1);

    if (attributes) {
        attributes_ = *attributes;

        for (auto& name_to_attr : attributes_) {
            if (name_to_attr.second.has_g()) {
                CreateSubgraph(name_to_attr.first);
            }
        }
    }
}

Node::Definitions& Node::MutableDefinitions() noexcept {
    // someone fetching these is going to change something
    graph_->SetGraphResolveNeeded();
    graph_->SetGraphProtoSyncNeeded();
    return definitions_;
}

Node::Relationships& Node::MutableRelationships() noexcept {
    // someone fetching these is going to change something
    graph_->SetGraphResolveNeeded();
    graph_->SetGraphProtoSyncNeeded();
    return relationships_;
}

void Node::CreateSubgraph(const std::string& attr_name) {
    auto attr = attributes_.find(attr_name);

    if (attr != attributes_.cend() && attr->second.has_g()) {
        GraphProto& mutable_graph = *attr->second.mutable_g();
        std::unique_ptr<Graph> subgraph{new Graph(*graph_, mutable_graph)};
        attr_to_subgraph_map_[attr_name] = subgraph.get();
        subgraphs_.push_back(std::move(subgraph));
    }
}

void Node::AddAttribute(const std::string& attr_name, const AttributeProto& value) {
    graph_->SetGraphResolveNeeded();
    graph_->SetGraphProtoSyncNeeded();
    attributes_[attr_name] = value;
}

#define ADD_BASIC_ATTR_IMPL(type, enumType, field)                                                                     \
    void Node::AddAttribute(const std::string& attr_name, const type& value) {                                         \
        graph_->SetGraphResolveNeeded();                                                                               \
        graph_->SetGraphProtoSyncNeeded();                                                                             \
        AttributeProto a;                                                                                              \
        a.set_name(attr_name);                                                                                         \
        a.set_type(enumType);                                                                                          \
        a.set_##field(value);                                                                                          \
        attributes_[attr_name] = a;                                                                                    \
    };

#define ADD_ATTR_IMPL(type, enumType, field)                                                                           \
    void Node::AddAttribute(const std::string& attr_name, const type& value) {                                         \
        graph_->SetGraphResolveNeeded();                                                                               \
        graph_->SetGraphProtoSyncNeeded();                                                                             \
        AttributeProto a;                                                                                              \
        a.set_name(attr_name);                                                                                         \
        a.set_type(enumType);                                                                                          \
        *(a.mutable_##field()) = value;                                                                                \
        attributes_[attr_name] = a;                                                                                    \
    };

#define ADD_LIST_ATTR_IMPL(type, enumType, field)                                                                      \
    void Node::AddAttribute(const std::string& attr_name, const std::vector<type>& values) {                           \
        graph_->SetGraphResolveNeeded();                                                                               \
        graph_->SetGraphProtoSyncNeeded();                                                                             \
        AttributeProto a;                                                                                              \
        a.set_name(attr_name);                                                                                         \
        a.set_type(enumType);                                                                                          \
        for (const auto& val : values) {                                                                               \
            *(a.mutable_##field()->Add()) = val;                                                                       \
        }                                                                                                              \
        attributes_[attr_name] = a;                                                                                    \
    };

void Node::AddAttribute(const std::string& attr_name, const GraphProto& value) {
    graph_->SetGraphResolveNeeded();
    graph_->SetGraphProtoSyncNeeded();
    AttributeProto a;
    a.set_name(attr_name);
    a.set_type(AttributeProto_AttributeType::AttributeProto_AttributeType_GRAPH);
    *a.mutable_g() = value;
    attributes_[attr_name] = a;

    CreateSubgraph(attr_name);
};

ADD_BASIC_ATTR_IMPL(float, AttributeProto_AttributeType::AttributeProto_AttributeType_FLOAT, f)
ADD_BASIC_ATTR_IMPL(int64_t, AttributeProto_AttributeType::AttributeProto_AttributeType_INT, i)
ADD_BASIC_ATTR_IMPL(std::string, AttributeProto_AttributeType::AttributeProto_AttributeType_STRING, s)
ADD_ATTR_IMPL(TensorProto, AttributeProto_AttributeType::AttributeProto_AttributeType_TENSOR, t)
ADD_LIST_ATTR_IMPL(float, AttributeProto_AttributeType::AttributeProto_AttributeType_FLOATS, floats)
ADD_LIST_ATTR_IMPL(int64_t, AttributeProto_AttributeType::AttributeProto_AttributeType_INTS, ints)
ADD_LIST_ATTR_IMPL(std::string, AttributeProto_AttributeType::AttributeProto_AttributeType_STRINGS, strings)
ADD_LIST_ATTR_IMPL(TensorProto, AttributeProto_AttributeType::AttributeProto_AttributeType_TENSORS, tensors)
ADD_LIST_ATTR_IMPL(GraphProto, AttributeProto_AttributeType::AttributeProto_AttributeType_GRAPHS, graphs)

bool Node::ClearAttribute(const std::string& attr_name) {
    graph_->SetGraphResolveNeeded();
    graph_->SetGraphProtoSyncNeeded();
    return attributes_.erase(attr_name) > 0;
}

Status Node::UpdateInputArgCount() {
    // The node refers to a primitive operator.
    // Infer and verify node input arg type information.
    int total_arg_count =
        std::accumulate(definitions_.input_arg_count.cbegin(), definitions_.input_arg_count.cend(), 0);

    if (total_arg_count < 0 || static_cast<size_t>(total_arg_count) != definitions_.input_defs.size()) {
        return FASTFLOW_MAKE_STATUS(
            FASTFLOW, FAIL, "The sum of input arg count is not equal to size of input defs in node (", name_, ")");
    }

    // op_ is always valid when this is called
    const ONNX_NAMESPACE::OpSchema& op = *Op();

    // Verify size of node arg count is same as input number in
    // operator definition.
    if (op.inputs().size() != definitions_.input_arg_count.size()) {
        // Adjust input arg count array with op definition
        // The adjustment will work as below,
        // In total, there're <total_arg_count> inputs, which
        // will be split as <1, 1, 1, 1, ... 1, x> or
        // <1, 1, 1, 1, ...1, 0, 0, ...0>. The final input
        // arg count array's element number will be the same
        // as op definition, and the sum of all elements will
        // be equal to <total_arg_count>.
        auto& input_arg_count = definitions_.input_arg_count;
        input_arg_count.clear();
        size_t m = 0;
        auto arg_count_left = total_arg_count;

        if (!op.inputs().empty()) {
            for (; m < op.inputs().size() - 1; ++m) {
                if (arg_count_left > 0) {
                    input_arg_count.push_back(1);
                    arg_count_left--;
                } else {
                    input_arg_count.push_back(0);
                }
            }
        }

        // Set the arg count for the last input formal parameter.
        // NOTE: in the case that there's no .input(...) defined
        // in op schema, all input args will be fed as one input
        // of the operator.
        input_arg_count.push_back(arg_count_left);

        graph_->SetGraphResolveNeeded();
        graph_->SetGraphProtoSyncNeeded();
    }

    return Status::OK();
}

const NodeAttributes& Node::GetAttributes() const noexcept {
    return attributes_;
}

Graph* Node::GetMutableGraphAttribute(const std::string& attr_name) {
    Graph* subgraph = nullptr;

    const auto& entry = attr_to_subgraph_map_.find(attr_name);
    if (entry != attr_to_subgraph_map_.cend()) {
        subgraph = entry->second;
    }

    return subgraph;
}

const Graph* Node::GetGraphAttribute(const std::string& attr_name) const {
    return const_cast<Node*>(this)->GetMutableGraphAttribute(attr_name);
}

void Node::ForEachDef(std::function<void(const fastflow::NodeArg&, bool is_input)> func) const {
    for (const auto* arg : InputDefs()) {
        if (arg->Exists())
            func(*arg, true);
    }

    for (const auto* arg : ImplicitInputDefs()) {
        if (arg->Exists())
            func(*arg, true);
    }

    for (const auto* arg : OutputDefs()) {
        if (arg->Exists())
            func(*arg, false);
    }
};

void Node::ReplaceDefs(const std::map<const fastflow::NodeArg*, fastflow::NodeArg*>& replacements) {
    std::vector<std::vector<NodeArg*>*> all_defs = {&definitions_.input_defs, &definitions_.output_defs};

    for (auto pair : replacements)
        for (auto* defs : all_defs)
            for (auto& def : *defs)
                if (def == pair.first)
                    def = pair.second;
}

}  // namespace fastflow
